"""
@Desc: 统计套利（Pairs Trading）/ 跨市场套利 / 时间套利
统计套利是一种基于历史统计关系的套利策略，通常用于两只相关性较高的股票或资产。
策略逻辑
选择资产对：选择两只历史价格高度相关的资产（如两只同行业股票）。
计算价差：计算两只资产的价格差或比率。
构建交易信号：
当价差偏离历史均值时，买入低估资产，卖出高估资产。
当价差回归均值时，平仓获利。

@Auth: meihongliang-m2
@Date: 2025/3/18-16:05
"""

import matplotlib.pyplot as plt
import yfinance as yf


# 1. 获取数据
def get_data(tickers, start_date, end_date):
    data = yf.download(tickers, start=start_date, end=end_date)['Adj Close']
    return data


# 2. 计算价差和Z-Score
def calculate_spread(data, ticker1, ticker2):
    data['Spread'] = data[ticker1] - data[ticker2]  # 计算价差
    data['Spread_Mean'] = data['Spread'].rolling(window=30).mean()  # 价差均值
    data['Spread_Std'] = data['Spread'].rolling(window=30).std()  # 价差标准差
    data['Z-Score'] = (data['Spread'] - data['Spread_Mean']) / data['Spread_Std']  # Z-Score
    return data


# 3. 生成交易信号
def generate_signals(data):
    data['Signal'] = 0
    data['Signal'][data['Z-Score'] > 1] = -1  # 价差过高，卖出ticker1，买入ticker2
    data['Signal'][data['Z-Score'] < -1] = 1  # 价差过低，买入ticker1，卖出ticker2
    data['Position'] = data['Signal'].diff()  # 计算仓位变化
    return data


# 4. 回测策略
def backtest(data, ticker1, ticker2):
    initial_capital = 10000  # 初始资金
    position_ticker1 = 0  # ticker1持仓
    position_ticker2 = 0  # ticker2持仓
    portfolio_value = []  # 记录每日投资组合价值

    for i in range(len(data)):
        if data['Position'][i] == 1:  # 买入ticker1，卖出ticker2
            position_ticker1 = initial_capital // (2 * data[ticker1][i])  # 买入ticker1
            position_ticker2 = -initial_capital // (2 * data[ticker2][i])  # 卖出ticker2
        elif data['Position'][i] == -1:  # 卖出ticker1，买入ticker2
            position_ticker1 = -initial_capital // (2 * data[ticker1][i])  # 卖出ticker1
            position_ticker2 = initial_capital // (2 * data[ticker2][i])  # 买入ticker2

        # 计算每日组合价值
        portfolio_value.append(position_ticker1 * data[ticker1][i] + position_ticker2 * data[ticker2][i])

    data['Portfolio_Value'] = portfolio_value
    return data


# 5. 可视化结果
def plot_results(data, ticker1, ticker2):
    plt.figure(figsize=(14, 7))
    plt.plot(data['Spread'], label='Spread', alpha=0.75)
    plt.plot(data['Spread_Mean'], label='Spread Mean', alpha=0.75)
    plt.scatter(data.index, data['Spread'], c=data['Position'], cmap='viridis', label='Buy/Sell Signal', marker='^',
                s=100)
    plt.title(f'Pairs Trading Strategy: {ticker1} vs {ticker2}')
    plt.legend()
    plt.show()


# 主程序
if __name__ == "__main__":
    # 参数设置
    ticker1 = 'AAPL'  # 股票1
    ticker2 = 'MSFT'  # 股票2
    start_date = '2020-01-01'
    end_date = '2023-01-01'

    # 获取数据
    data = get_data([ticker1, ticker2], start_date, end_date)
    # 计算价差和Z-Score
    data = calculate_spread(data, ticker1, ticker2)
    # 生成信号
    data = generate_signals(data)
    # 回测策略
    data = backtest(data, ticker1, ticker2)
    # 可视化结果
    plot_results(data, ticker1, ticker2)
    # 打印最终收益
    final_portfolio_value = data['Portfolio_Value'].iloc[-1]
    print(f"Final Portfolio Value: {final_portfolio_value:.2f}")
